{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# importing libs\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import keras\n",
    "from keras.layers import Input, Dense, GaussianNoise,Lambda,Dropout\n",
    "from keras.models import Model\n",
    "from keras import regularizers\n",
    "from keras.layers.normalization import BatchNormalization\n",
    "from keras.optimizers import Adam,SGD\n",
    "from keras import backend as K"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# for reproducing reslut\n",
    "from numpy.random import seed\n",
    "seed(1)\n",
    "from tensorflow import set_random_seed\n",
    "set_random_seed(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining parameters\n",
    "# define (n,k) here for (n,k) autoencoder\n",
    "# n = n_channel \n",
    "# k = log2(M)  ==> so for (7,4) autoencoder n_channel = 7 and M = 2^4 = 16 \n",
    "M = 4\n",
    "k = np.log2(M)\n",
    "k = int(k)\n",
    "n_channel = 2\n",
    "R = k/n_channel\n",
    "print ('M:',M,'k:',k,'n:',n_channel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#generating data of size N\n",
    "N = 8000\n",
    "label = np.random.randint(M,size=N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# creating one hot encoded vectors\n",
    "data = []\n",
    "for i in label:\n",
    "    temp = np.zeros(M)\n",
    "    temp[i] = 1\n",
    "    data.append(temp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# checking data shape\n",
    "data = np.array(data)\n",
    "print (data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# checking generated data with it's label\n",
    "temp_check = [17,23,45,67,89,96,72,250,350]\n",
    "for i in temp_check:\n",
    "    print(label[i],data[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# defining autoencoder and it's layer\n",
    "input_signal = Input(shape=(M,))\n",
    "encoded = Dense(M, activation='relu')(input_signal)\n",
    "encoded1 = Dense(n_channel, activation='linear')(encoded)\n",
    "encoded2 = Lambda(lambda x: np.sqrt(n_channel)*K.l2_normalize(x,axis=1))(encoded1)\n",
    "\n",
    "EbNo_train = 5.01187 #  coverted 7 db of EbNo\n",
    "encoded3 = GaussianNoise(np.sqrt(1/(2*R*EbNo_train)))(encoded2)\n",
    "\n",
    "decoded = Dense(M, activation='relu')(encoded3)\n",
    "decoded1 = Dense(M, activation='softmax')(decoded)\n",
    "autoencoder = Model(input_signal, decoded1)\n",
    "adam = Adam(lr=0.01)\n",
    "autoencoder.compile(optimizer=adam, loss='categorical_crossentropy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# printing summary of layers and it's trainable parameters \n",
    "print (autoencoder.summary())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# for tensor board visualization\n",
    "#tbCallBack = keras.callbacks.TensorBoard(log_dir='./logs', histogram_freq=0, batch_size=32, write_graph=True, write_grads=True, write_images=False, embeddings_freq=0, embeddings_layer_names=None, embeddings_metadata=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# traning auto encoder\n",
    "autoencoder.fit(data, data,\n",
    "                epochs=45,\n",
    "                batch_size=32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# saving keras model\n",
    "from keras.models import load_model\n",
    "# if you want to save model then remove below comment\n",
    "# autoencoder.save('autoencoder_v_best.model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# making encoder from full autoencoder\n",
    "encoder = Model(input_signal, encoded2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# making decoder from full autoencoder\n",
    "encoded_input = Input(shape=(n_channel,))\n",
    "\n",
    "deco = autoencoder.layers[-2](encoded_input)\n",
    "deco = autoencoder.layers[-1](deco)\n",
    "decoder = Model(encoded_input, deco)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# generating data for checking BER\n",
    "# if you're not using t-sne for visulation than set N to 70,000 for better result \n",
    "# for t-sne use less N like N = 1500\n",
    "N = 50000\n",
    "test_label = np.random.randint(M,size=N)\n",
    "test_data = []\n",
    "\n",
    "for i in test_label:\n",
    "    temp = np.zeros(M)\n",
    "    temp[i] = 1\n",
    "    test_data.append(temp)\n",
    "    \n",
    "test_data = np.array(test_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# checking generated data\n",
    "temp_test = 6\n",
    "print (test_data[temp_test][test_label[temp_test]],test_label[temp_test])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for plotting learned consteallation diagram\n",
    "\n",
    "scatter_plot = []\n",
    "for i in range(0,M):\n",
    "    temp = np.zeros(M)\n",
    "    temp[i] = 1\n",
    "    scatter_plot.append(encoder.predict(np.expand_dims(temp,axis=0)))\n",
    "scatter_plot = np.array(scatter_plot)\n",
    "print (scatter_plot.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    " # use this function for ploting constellation for higher dimenson like 7-D for (7,4) autoencoder \n",
    "'''\n",
    "x_emb = encoder.predict(test_data)\n",
    "noise_std = np.sqrt(1/(2*R*EbNo_train))\n",
    "noise = noise_std * np.random.randn(N,n_channel)\n",
    "x_emb = x_emb + noise\n",
    "from sklearn.manifold import TSNE\n",
    "X_embedded = TSNE(learning_rate=700, n_components=2,n_iter=35000, random_state=0, perplexity=60).fit_transform(x_emb)\n",
    "print (X_embedded.shape)\n",
    "X_embedded = X_embedded / 7\n",
    "import matplotlib.pyplot as plt\n",
    "plt.scatter(X_embedded[:,0],X_embedded[:,1])\n",
    "#plt.axis((-2.5,2.5,-2.5,2.5)) \n",
    "plt.grid()\n",
    "plt.show()\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ploting constellation diagram\n",
    "import matplotlib.pyplot as plt\n",
    "scatter_plot = scatter_plot.reshape(M,2,1)\n",
    "plt.scatter(scatter_plot[:,0],scatter_plot[:,1])\n",
    "plt.axis((-2.5,2.5,-2.5,2.5))\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def frange(x, y, jump):\n",
    "  while x < y:\n",
    "    yield x\n",
    "    x += jump"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculating BER\n",
    "# this is optimized BER function so it can handle large number of N\n",
    "# previous code has another for loop which was making it slow\n",
    "EbNodB_range = list(frange(-4,8.5,0.5))\n",
    "ber = [None]*len(EbNodB_range)\n",
    "for n in range(0,len(EbNodB_range)):\n",
    "    EbNo=10.0**(EbNodB_range[n]/10.0)\n",
    "    noise_std = np.sqrt(1/(2*R*EbNo))\n",
    "    noise_mean = 0\n",
    "    no_errors = 0\n",
    "    nn = N\n",
    "    noise = noise_std * np.random.randn(nn,n_channel)\n",
    "    encoded_signal = encoder.predict(test_data) \n",
    "    final_signal = encoded_signal + noise\n",
    "    pred_final_signal =  decoder.predict(final_signal)\n",
    "    pred_output = np.argmax(pred_final_signal,axis=1)\n",
    "    no_errors = (pred_output != test_label)\n",
    "    no_errors =  no_errors.astype(int).sum()\n",
    "    ber[n] = no_errors / nn \n",
    "    print ('SNR:',EbNodB_range[n],'BER:',ber[n])\n",
    "    # use below line for generating matlab like matrix which can be copy and paste for plotting ber graph in matlab\n",
    "    #print(ber[n], \" \",end='')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ploting ber curve\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import interpolate\n",
    "plt.plot(EbNodB_range, ber, 'bo',label='Autoencoder(2,2)')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('SNR Range')\n",
    "plt.ylabel('Block Error Rate')\n",
    "plt.grid()\n",
    "plt.legend(loc='upper right',ncol = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for saving figure remove below comment\n",
    "#plt.savefig('AutoEncoder_2_2_constrained_BER_matplotlib')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
